"""
训练任务路由
"""
from typing import Optional, List
from fastapi import APIRouter, Depends, Query
from sqlalchemy.orm import Session
from database import get_db, TrainingTask
from schemas import (
    TrainingTaskCreate, TrainingTaskResponse, TrainingTaskListResponse,
    PaginationParams, TrainingTaskUpdate, TrainingTaskProgressUpdate, GenericResponse, ListResponse, DataResponse,
    TaskLogResponse, TaskLogRequest
)
from services import TrainingService
from exceptions import NotFoundException, ValidationException
from response_utils import success_response, create_pagination_info, list_response, error_response

router = APIRouter(prefix="/api/v1", tags=["训练任务"])

# 训练服务实例
training_service = TrainingService()

@router.post("/training-tasks", response_model=DataResponse[TrainingTaskResponse])
async def create_training_task(
    task_data: TrainingTaskCreate,
    db: Session = Depends(get_db)
):
    """创建训练任务"""
    task = training_service.create_training_task(db, task_data)
    return success_response(data=task, message="训练任务创建成功")

@router.get("/training-tasks", response_model=ListResponse[TrainingTaskResponse])
async def list_training_tasks(
    page: int = 1,
    page_size: int = 20,
    status: Optional[str] = None,
    model_type: Optional[str] = None,
    db: Session = Depends(get_db)
):
    """获取训练任务列表"""
    tasks, total_count = training_service.list_training_tasks(
        db, page=page, page_size=page_size, status=status, model_type=model_type
    )
    pagination_info = create_pagination_info(page, page_size, total_count)
    
    # 调用服务层的统计方法
    summary = training_service.get_training_tasks_summary(db, status, model_type)
    
    return list_response(
        items=tasks,
        total_count=total_count,
        pagination=pagination_info,
        summary=summary,
        message="查询成功"
    )

@router.get("/training-tasks/{task_id}", response_model=DataResponse[TrainingTaskResponse])
async def get_training_task(task_id: str, db: Session = Depends(get_db)):
    """获取训练任务详情"""
    task = training_service.get_training_task_by_id(db, task_id)
    return success_response(data=task, message="查询成功")

@router.get("/training-tasks/{task_id}/logs", response_model=DataResponse[TaskLogResponse])
async def get_training_task_logs(
    task_id: str,
    start_line: int = Query(1, description="起始行号，默认从第1行开始"),
    end_line: Optional[int] = Query(None, description="结束行号，为空则读取到文件末尾"),
    max_lines: int = Query(1000, description="最大返回行数，默认1000行"),
    level_filter: Optional[str] = Query(None, description="日志级别过滤：INFO, DEBUG, WARNING, ERROR")
):
    """获取训练任务日志"""
    logs = training_service.get_task_logs(
        task_id=task_id,
        start_line=start_line,
        end_line=end_line,
        max_lines=max_lines,
        level_filter=level_filter
    )
    return success_response(data=logs, message="获取日志成功")

@router.put("/training-tasks/{task_id}", response_model=DataResponse[TrainingTaskResponse])
async def update_training_task(
    task_id: str,
    task_update: TrainingTaskUpdate,
    db: Session = Depends(get_db)
):
    """更新训练任务"""
    updated_task = training_service.update_training_task(db, task_id, task_update)
    return success_response(data=updated_task, message="训练任务更新成功")

@router.patch("/training-tasks/{task_id}/progress", response_model=DataResponse[TrainingTaskResponse])
async def update_training_task_progress(
    task_id: str,
    progress_update: TrainingTaskProgressUpdate,
    db: Session = Depends(get_db)
):
    """更新训练任务进度"""
    updated_task = training_service.update_task_progress(
        db, task_id, progress_update.progress, progress_update.accuracy
    )
    return success_response(data=updated_task, message="训练任务进度更新成功")

@router.post("/training-tasks/{task_id}/cancel", response_model=GenericResponse)
async def cancel_training_task(task_id: str, db: Session = Depends(get_db)):
    """取消训练任务"""
    training_service.cancel_training_task(db, task_id)
    return success_response(message="训练任务已取消")

@router.delete("/training-tasks/{task_id}", response_model=GenericResponse)
async def delete_training_task(task_id: str, db: Session = Depends(get_db)):
    """删除训练任务和相关模型"""
    training_service.delete_training_task(db, task_id)
    return success_response(message="训练任务和相关模型已删除")

@router.post("/training-tasks/{task_id}/retry", response_model=DataResponse[TrainingTaskResponse])
async def retry_training_task(task_id: str, db: Session = Depends(get_db)):
    """重试训练任务"""
    retried_task = training_service.retry_training_task(db, task_id)
    return success_response(data=retried_task, message="训练任务重试已启动") 

